{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "# Based on this paper - https://arxiv.org/abs/1409.0473\n",
    "#!pip install torch torchtext sentencepiece datasets"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import torch\n",
    "from torch import nn\n",
    "import functorch\n",
    "import sys\n",
    "import os\n",
    "import math\n",
    "sys.path.append(os.path.abspath(\"../../data\"))\n",
    "\n",
    "DEVICE = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "BATCH_SIZE = 8\n",
    "SP_VOCAB_SIZE = 1000\n",
    "TRAIN_SIZE = 500"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from text_data import Opus100DatasetWrapper\n",
    "\n",
    "class Wrapper(Opus100DatasetWrapper):\n",
    "    split_lengths = [TRAIN_SIZE, math.floor(TRAIN_SIZE * .1), 100]\n",
    "    x_length = 40\n",
    "    target_length = 40\n",
    "\n",
    "wrapper = Wrapper(SP_VOCAB_SIZE, DEVICE)\n",
    "\n",
    "datasets = wrapper.generate_datasets(BATCH_SIZE)\n",
    "train = datasets[\"train\"]\n",
    "valid = datasets[\"validation\"]"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class GRUCell(nn.Module):\n",
    "    def __init__(self, input_units, hidden_units, output_units):\n",
    "        super(GRUCell, self).__init__()\n",
    "        self.input_units = input_units\n",
    "        self.hidden_units = hidden_units\n",
    "        self.output_units = output_units\n",
    "\n",
    "        k = math.sqrt(1/hidden_units)\n",
    "        self.input_weights = nn.Parameter(torch.rand(3, input_units, hidden_units) * 2 * k - k)\n",
    "        self.input_biases = nn.Parameter(torch.rand(3, 1, hidden_units) * 2 * k - k)\n",
    "\n",
    "        self.hidden_weights = nn.Parameter(torch.rand(3, hidden_units, hidden_units) * 2 * k - k)\n",
    "        self.hidden_biases = nn.Parameter(torch.rand(3, 1, hidden_units) * 2 * k - k)\n",
    "\n",
    "    def forward(self, x, prev_hidden):\n",
    "        # Compute the regular RNN forward pass\n",
    "        # Compute update and reset gates for GRU\n",
    "        reset_gate = torch.sigmoid(x @ self.input_weights[0,] + self.input_biases[0,] + prev_hidden @ self.hidden_weights[0,] + self.hidden_biases[0,])\n",
    "        update_gate = torch.sigmoid(x @ self.input_weights[1,] + self.input_biases[1,] + prev_hidden @ self.hidden_weights[1,] + self.hidden_biases[1,])\n",
    "        new_gate = torch.tanh(x @ self.input_weights[2,] + self.input_biases[2,] + torch.mul(reset_gate, prev_hidden @ self.hidden_weights[2,] + self.hidden_biases[2,]))\n",
    "\n",
    "        hidden_x = torch.mul((1 - update_gate), new_gate) + torch.mul(update_gate, new_gate)\n",
    "        return hidden_x"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "class EncoderDecoder(nn.Module):\n",
    "    def __init__(self, in_sequence_len, out_sequence_len, embedding_len, hidden_units=512, layers=2, bidirectional=True):\n",
    "        super(EncoderDecoder, self).__init__()\n",
    "        self.in_sequence_len = in_sequence_len\n",
    "        self.out_sequence_len = out_sequence_len\n",
    "        self.hidden_units = hidden_units\n",
    "        # Bidirectional rnn needs 2x the params\n",
    "        if bidirectional:\n",
    "            self.context_units = self.hidden_units * 2\n",
    "        else:\n",
    "            self.context_units = self.hidden_units\n",
    "\n",
    "        self.bidirectional = bidirectional\n",
    "        self.embedding_len = embedding_len\n",
    "        self.layers = layers\n",
    "\n",
    "        self.embedding = nn.Embedding(embedding_len, hidden_units)\n",
    "        self.fwd_encoders = nn.ModuleList([GRUCell(input_units=hidden_units, hidden_units=hidden_units, output_units=hidden_units) for _ in range(layers)])\n",
    "        self.bwd_encoders = nn.ModuleList([GRUCell(input_units=hidden_units, hidden_units=hidden_units, output_units=hidden_units) for _ in range(layers)])\n",
    "        self.decoders = nn.ModuleList([GRUCell(input_units=self.context_units + hidden_units, hidden_units=hidden_units, output_units=hidden_units) for _ in range(layers)])\n",
    "\n",
    "        self.linear = nn.Linear(in_features=hidden_units, out_features=embedding_len)\n",
    "\n",
    "        k = math.sqrt(1/hidden_units)\n",
    "        self.hidden_attention_weight = nn.Parameter(torch.rand(hidden_units, hidden_units) * 2 * k - k)\n",
    "        self.context_attention_weight = nn.Parameter(torch.rand(self.context_units, hidden_units) * 2 * k - k)\n",
    "        self.attention_weight = nn.Parameter(torch.rand(1, hidden_units) * 2 * k - k)\n",
    "        self.batched_diag = functorch.vmap(torch.diag)\n",
    "\n",
    "    def attention(self, context, prev_hidden, batch_size):\n",
    "        # Swap axes so the first dimension of context_attn is batch\n",
    "        context_attn = torch.bmm(context.swapaxes(0,1), self.context_attention_weight.unsqueeze(0).expand(batch_size,-1,-1))\n",
    "        # Swap back since prev_hidden is by batch.  This makes the first dim of cross sequence\n",
    "        cross = torch.tanh(context_attn.swapaxes(0,1) + prev_hidden @ self.hidden_attention_weight)\n",
    "        # This will be of dimension batch, sequence_length, 1\n",
    "        attention = torch.bmm(cross.swapaxes(0,1), self.attention_weight.T.unsqueeze(0).expand(batch_size, -1, -1))\n",
    "        # Drop the last singleton dimension\n",
    "        attention = attention.squeeze(2)\n",
    "        # Softmax the predictions across each batch\n",
    "        probs = torch.softmax(attention, 1)\n",
    "        diagonalized_probs = self.batched_diag(probs)\n",
    "        positional_contexts = torch.sum(torch.bmm(diagonalized_probs, context.swapaxes(0,1)), dim=1).reshape(batch_size, self.context_units)\n",
    "        return positional_contexts\n",
    "\n",
    "    def generate_context(self, embedded, batch_size, encoders):\n",
    "        # Encode the input sequence\n",
    "        # Both tensors will have sequence then batch\n",
    "        enc_hiddens = torch.zeros((1, self.layers, batch_size, self.hidden_units), device=DEVICE)\n",
    "        for j in range(self.in_sequence_len):\n",
    "            seq_enc_hiddens = embedded[j,:].unsqueeze(0)\n",
    "            for i in range(self.layers):\n",
    "                hidden = encoders[i](seq_enc_hiddens[i,], enc_hiddens[j,i])\n",
    "                # Add first sequence axis\n",
    "                hidden = hidden.unsqueeze(0)\n",
    "                seq_enc_hiddens = torch.cat((seq_enc_hiddens, hidden), dim=0)\n",
    "\n",
    "            enc_hiddens = torch.cat((enc_hiddens, seq_enc_hiddens[1:].unsqueeze(0)), dim=0)\n",
    "\n",
    "        # Decode to the output sequence\n",
    "        # Pass in context\n",
    "        context = enc_hiddens[1:,-1,:,:]\n",
    "\n",
    "        return context\n",
    "\n",
    "    def forward(self, x, y):\n",
    "        batch_size = x.shape[0]\n",
    "        # Move batch to the second dimension, so sequence comes first\n",
    "        y = y.swapaxes(0,1)\n",
    "        # Embed the input sequence to reduce dimensionality\n",
    "        embedded = self.embedding(x).swapaxes(0,1)\n",
    "\n",
    "        # Generate forward context (go through input sequence from the start)\n",
    "        forward_context = self.generate_context(embedded, batch_size, self.fwd_encoders)\n",
    "\n",
    "        if self.bidirectional:\n",
    "            # Generate backward context (go through input sequence from the end)\n",
    "            backward_context = torch.flip(self.generate_context(torch.flip(embedded, [0]), batch_size, self.bwd_encoders), [0])\n",
    "            context = torch.cat((forward_context, backward_context), dim=2)\n",
    "        else:\n",
    "            context = forward_context\n",
    "        # Both tensors will have the first dimension be the sequence\n",
    "        dec_hiddens = torch.zeros(1, self.layers, batch_size, self.hidden_units, device=DEVICE)\n",
    "        outputs = torch.zeros(1, batch_size, self.embedding_len, device=DEVICE)\n",
    "        for j in range(self.out_sequence_len):\n",
    "            # Use either the actual previous y (from the input), or the generated y if the input sequence is shorter than the generation steps.\n",
    "            if y.shape[0] > j:\n",
    "                prev_y = y[j,:]\n",
    "            else:\n",
    "                prev_y = outputs[j,:,:]\n",
    "                prev_y = prev_y.argmax(dim=1).int()\n",
    "\n",
    "            # Run embedding over previous y state\n",
    "            prev_y = self.embedding(prev_y)\n",
    "            seq_dec_hiddens = prev_y.unsqueeze(0)\n",
    "            for i in range(self.layers):\n",
    "                positional_context = self.attention(context, dec_hiddens[j,i,], batch_size)\n",
    "                hidden = self.decoders[i](torch.cat((seq_dec_hiddens[i,], positional_context), dim=1), dec_hiddens[j,i,],)\n",
    "                # Add first sequence axis\n",
    "                hidden = hidden.unsqueeze(0)\n",
    "                seq_dec_hiddens = torch.cat((seq_dec_hiddens, hidden), dim=0)\n",
    "\n",
    "            # Swap sequence and batch axes to apply linear transform, then swap back\n",
    "            prev_output = self.linear(seq_dec_hiddens[-1]).unsqueeze(0)\n",
    "            outputs = torch.cat((outputs, prev_output), dim=0)\n",
    "            dec_hiddens = torch.cat((dec_hiddens, seq_dec_hiddens[1:].unsqueeze(0)), dim=0)\n",
    "\n",
    "        # Move batch back to axis 0\n",
    "        out_hiddens = dec_hiddens[1:,-1,:,:].swapaxes(0,1)\n",
    "        out_output = outputs[1:,].swapaxes(0,1)\n",
    "        return out_output, out_hiddens\n",
    "\n",
    "def generate(sequence, prev_target, target, wrapper):\n",
    "    pred, _ = model(sequence, prev_target[:,0].unsqueeze(1))\n",
    "    prompts = wrapper.decode_batch(sequence.cpu())\n",
    "    texts = wrapper.decode_batch(torch.argmax(pred, dim=2).cpu())\n",
    "    correct_texts = wrapper.decode_batch(target.cpu())\n",
    "\n",
    "    displays = []\n",
    "    for p, t, ct in zip(prompts, texts, correct_texts):\n",
    "        displays.append(f\"{p} | {ct} | {t}\")\n",
    "    return displays"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "from tqdm.auto import tqdm\n",
    "model = EncoderDecoder(wrapper.x_length, wrapper.y_length, hidden_units=512, layers=1, embedding_len=wrapper.vocab_size).to(DEVICE)\n",
    "loss_fn = nn.CrossEntropyLoss(ignore_index=wrapper.pad_token)\n",
    "optimizer = torch.optim.AdamW(model.parameters(), lr=1e-4)"
   ],
   "metadata": {
    "collapsed": false
   }
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "outputs": [],
   "source": [
    "EPOCHS = 1000\n",
    "DISPLAY_BATCHES = 2\n",
    "\n",
    "for epoch in range(EPOCHS):\n",
    "    # Run over the training examples\n",
    "    train_loss = 0\n",
    "    match_pct = 0\n",
    "    for batch, (sequence, target, prev_target) in tqdm(enumerate(train)):\n",
    "        optimizer.zero_grad()\n",
    "        forced_target = prev_target\n",
    "        # Alternate use of teacher forcing vs feeding back own inputs\n",
    "        if np.random.randint(2) == 0:\n",
    "            forced_target = prev_target[:,0].unsqueeze(1)\n",
    "        pred, hidden = model(sequence, forced_target)\n",
    "\n",
    "        # Need to reshape pred to be batch * sequence, embedding_len to be compatible\n",
    "        # Similar reshape with target to be batch * sequence vector of class indices\n",
    "        loss = loss_fn(pred.reshape(-1, pred.shape[-1]), target.view(-1))\n",
    "        loss.backward()\n",
    "        optimizer.step()\n",
    "        train_loss += loss.item()\n",
    "        match_pct += torch.sum(target == torch.argmax(pred, 2)) / (wrapper.y_length * BATCH_SIZE)\n",
    "\n",
    "    with torch.no_grad():\n",
    "        if epoch % 10 == 0:\n",
    "            # Show text generated from training prompt as an example\n",
    "            # Don't feed in all of the train y sequences, just the first token\n",
    "            # The other y tokens will be predicted by the model and fed back in\n",
    "            print(f\"Epoch {epoch} train loss: {train_loss / len(train)} match_pct: {match_pct / len(train)}\")\n",
    "            sents = generate(sequence[:DISPLAY_BATCHES], prev_target[:DISPLAY_BATCHES], target[:DISPLAY_BATCHES], wrapper)\n",
    "            for sent in sents:\n",
    "                print(sent)\n",
    "\n",
    "            # Compute validation loss.  Unless you have a lot of training data, the validation loss won't decrease.\n",
    "            valid_loss = 0\n",
    "            for batch, (sequence, target, prev_target) in enumerate(valid):\n",
    "                # Only feed in the first token of the actual target\n",
    "                pred, hidden = model(sequence, prev_target[:,0].unsqueeze(1))\n",
    "                loss = loss_fn(pred.reshape(-1, pred.shape[-1]), target.view(-1))\n",
    "                valid_loss += loss.item()\n",
    "\n",
    "            print(f\"Valid loss: {valid_loss / len(valid)}\")\n",
    "            sents = generate(sequence[:DISPLAY_BATCHES], prev_target[:DISPLAY_BATCHES], target[:DISPLAY_BATCHES], wrapper)\n",
    "            for sent in sents:\n",
    "                print(sent)"
   ],
   "metadata": {
    "collapsed": false
   }
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 0
}
